from flask import Flask, Response, stream_with_context, request, jsonify
from flask_cors import CORS
import requests
import json
import os

app = Flask(__name__)
# 配置跨域，允许所有来源的请求，并支持所有HTTP方法和 headers
CORS(
    app,
    resources={
        r"/*": {
            "origins": "*",
            "methods": ["GET", "POST", "OPTIONS"],
            "allow_headers": "*",
        }
    },
)

# 配置（建议通过环境变量注入）
SILICONFLOW_API_KEY = os.environ.get(
    "SILICONFLOW_API_KEY", "sk-svuwiezizbxyflaqnhapzgijjdmeiaotknizeezdwzesoruw"
)
SILICONFLOW_API_URL = "https://api.siliconflow.cn/v1/chat/completions"
LLM_MODEL = "deepseek-ai/DeepSeek-V3"  # 替换为硅基流动支持的模型（如 siliconflow-pro）


@app.route("/stream-completion", methods=["GET"])  # 修改为GET方法以适配EventSource
def stream_completion():
    """
    流式调用硅基流动API，返回SSE格式响应
    """
    # 从URL参数获取prompt
    prompt = request.args.get("prompt", "")
    if not prompt:
        return Response(
            response=json.dumps({"error": "缺少prompt参数"}),
            status=400,
            mimetype="application/json",
        )

    # 构造硅基流动请求体
    llm_payload = {
        "model": LLM_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "stream": True,  # 开启流式响应
    }
    headers = {
        "Authorization": f"Bearer {SILICONFLOW_API_KEY}",
        "Content-Type": "application/json",
    }

    def generate():
        """生成器：逐段处理硅基流动的响应"""
        try:
            # 发起流式请求
            resp = requests.post(
                SILICONFLOW_API_URL,
                headers=headers,
                json=llm_payload,
                stream=True,
                timeout=60,
            )
            resp.raise_for_status()  # 捕获HTTP错误
        except requests.exceptions.RequestException as e:
            # 异常时返回错误信息
            yield f"data: {json.dumps({'error': f'调用硅基流动失败: {str(e)}'})}\n\n"
            return

        # 逐行解析流式响应
        for line in resp.iter_lines():
            if not line:
                continue
            # 硅基流动的流式响应可能以"data:"开头，需要处理
            line_str = line.decode("utf-8").strip()
            if line_str.startswith("data: "):
                line_str = line_str[6:]  # 去掉"data: "前缀
            if line_str == "[DONE]":  # 结束标记
                break
            try:
                chunk = json.loads(line_str)
                # 解析文本内容
                choices = chunk.get("choices", [])
                if not choices:
                    continue
                delta = choices[0].get("delta", {})
                content = delta.get("content")
                if content:  # 仅当有新文本时发送
                    yield f"data: {json.dumps({'text': content})}\n\n"
            except json.JSONDecodeError as e:
                app.logger.error(f"解析响应失败: {line_str}, 错误: {str(e)}")
                continue

    # 返回流式响应（SSE格式）
    return Response(
        stream_with_context(generate()),
        mimetype="text/event-stream",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
    )


# 处理预检请求
@app.before_request
def handle_options():
    if request.method == "OPTIONS":
        response = jsonify({"status": "OK"})
        response.headers["Access-Control-Allow-Origin"] = "*"
        response.headers["Access-Control-Allow-Methods"] = "GET, POST, OPTIONS"
        response.headers["Access-Control-Allow-Headers"] = "*"
        return response, 200


from flask import render_template


@app.route("/")
def index():
    return render_template("index.html")


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=5000)
